import torch.nn as nn

class FCNet(nn.Module):
    def __init__(self,input_size,hidden_size1,hidden_size2,output_size):
        super(FCNet,self).__init__()
        self.fc1 = nn.Linear(input_size,hidden_size1)
        self.relu1 = nn.ReLU()
        self.fc2 = nn.Linear(hidden_size1,hidden_size2)
        self.relu2 = nn.ReLU()
        self.fc3 = nn.Linear(hidden_size2,output_size)
    def forward(self,x):
        x = self.fc1(x)
        x = self.relu1(x)
        x = self.fc2(x)
        x = self.relu2(x)
        x = self.fc3(x)
        return x
    
class CNNet(nn.Module):
    def __init__(self, input_size, output_size):
        super(CNNet, self).__init__()
        self.block1 = nn.Sequential(
            nn.Conv2d(1, 10, 5),
            nn.MaxPool2d(2),
            nn.ReLU(True),
            nn.BatchNorm2d(10),
        )
        self.block2 = nn.Sequential(
            nn.Conv2d(10, 20, 5),
            nn.MaxPool2d(2),
            nn.ReLU(True),
            nn.BatchNorm2d(20),
        )
        self.fc = nn.Sequential(
            nn.Flatten(),
            nn.Linear(320, output_size)
        )
    def forward(self, x):
        x = self.block1(x)
        x = self.block2(x)
        x = self.fc(x)
        return x
